package com.zcxy360.sandtable.websocket;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

public class WebSocketServer extends TextWebSocketHandler {

	private final Map<WebSocketSession, String> connections = new ConcurrentHashMap<>();
	private final Map<String, TextMessage> messageCache = new ConcurrentHashMap<>();

	/**
	 * 建立连接
	 */
	@Override
	public void afterConnectionEstablished(WebSocketSession session) throws Exception {
		String uri = session.getUri().toString();
		String teamId = uri.substring(uri.lastIndexOf("/") + 1);
		connections.put(session, teamId);

		TextMessage message = messageCache.get(teamId);
		if (message != null) {
			session.sendMessage(message);
		}
	}

	/**
	 * 断开连接
	 */
	@Override
	public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
		connections.remove(session);
	}

	/**
	 * 处理消息
	 */
	@Override
	protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
		broadcast(session, message);
	}

	private void broadcast(WebSocketSession session, TextMessage message) {
		String teamId = connections.get(session);
		for (WebSocketSession s : connections.keySet()) {
			if (connections.get(s).equals(teamId) && !s.equals(session)) {
				try {
					s.sendMessage(message);
				} catch (Exception e) {
					connections.remove(s);
					try {
						s.close();
					} catch (IOException e1) {
						e1.printStackTrace();
					}
				}
			}
		}
		messageCache.put(teamId, message);
	}
}